# -*- coding: utf-8 -*-
# Author   : ZhangQing
# Time     : 2025-07-15 23:08
# File     : data_sources.py
# Project  : dynamic-portfolio-optimizer
# Desc     : 数据源配置管理

import os
from dataclasses import dataclass
from typing import Dict, Optional


@dataclass
class DataSourceConfig:
    """数据源配置"""
    name: str
    api_key: Optional[str] = None
    base_url: Optional[str] = None
    rate_limit: int = 5  # 每秒请求数
    timeout: int = 30
    enabled: bool = True


class DataSourceManager:
    """数据源管理器"""

    def __init__(self):
        self.configs = {
            'yfinance': DataSourceConfig(
                name='Yahoo Finance',
                rate_limit=10,
                enabled=True
            ),
            'alpha_vantage': DataSourceConfig(
                name='Alpha Vantage',
                api_key=os.getenv('ALPHA_VANTAGE_API_KEY'),
                base_url='https://www.alphavantage.co',
                rate_limit=5,
                enabled=bool(os.getenv('ALPHA_VANTAGE_API_KEY'))
            ),
            'nasdaq_data_link': DataSourceConfig(
                name='Nasdaq Data Link',
                api_key=os.getenv('NASDAQ_DATA_LINK_API_KEY'),
                base_url='https://data.nasdaq.com',
                rate_limit=10,
                enabled=bool(os.getenv('NASDAQ_DATA_LINK_API_KEY'))
            ),
            'twelve_data': DataSourceConfig(
                name='Twelve Data',
                api_key=os.getenv('TWELVE_DATA_API_KEY'),
                base_url='https://api.twelvedata.com',
                rate_limit=8,
                enabled=bool(os.getenv('TWELVE_DATA_API_KEY'))
            ),
            'polygon': DataSourceConfig(
                name='Polygon.io',
                api_key=os.getenv('POLYGON_API_KEY'),
                base_url='https://api.polygon.io',
                rate_limit=5,
                enabled=bool(os.getenv('POLYGON_API_KEY'))
            ),
            'finnhub': DataSourceConfig(
                name='Finnhub',
                api_key=os.getenv('FINNHUB_API_KEY'),
                base_url='https://finnhub.io/api/v1',
                rate_limit=60,
                enabled=bool(os.getenv('FINNHUB_API_KEY'))
            ),
            'alpaca': DataSourceConfig(
                name='Alpaca',
                api_key=os.getenv('ALPACA_API_KEY'),
                base_url='https://data.alpaca.markets',
                rate_limit=200,
                enabled=bool(os.getenv('ALPACA_API_KEY'))
            ),
            'marketstack': DataSourceConfig(
                name='Marketstack',
                api_key=os.getenv('MARKETSTACK_API_KEY'),
                base_url='http://api.marketstack.com',
                rate_limit=10,
                enabled=bool(os.getenv('MARKETSTACK_API_KEY'))
            ),
            'tradier': DataSourceConfig(
                name='Tradier',
                api_key=os.getenv('TRADIER_API_KEY'),
                base_url='https://api.tradier.com',
                rate_limit=120,
                enabled=bool(os.getenv('TRADIER_API_KEY'))
            ),
            'ib_api': DataSourceConfig(
                name='Interactive Brokers',
                rate_limit=50,
                enabled=bool(os.getenv('IB_GATEWAY_ENABLED'))
            )
        }

    def get_enabled_sources(self) -> Dict[str, DataSourceConfig]:
        """获取已启用的数据源"""
        return {k: v for k, v in self.configs.items() if v.enabled}

    def get_source_config(self, source_name: str) -> Optional[DataSourceConfig]:
        """获取特定数据源配置"""
        return self.configs.get(source_name)


if __name__ == '__main__':
    DataSourceManager()